import torch

from problems.DTLZ import DTLZ


class DTLZ2(DTLZ):
    def __init__(self, var_dim: int, obj_dim: int, max_fun_eval, kwargs: dict):
        super().__init__(var_dim, obj_dim, max_fun_eval, kwargs)

    def eval_value(self, dec):
        g = torch.sum((dec[:, self.obj_dim - 1:] - 0.5) ** 2, dim=1)
        ones_vector = torch.ones((dec.size(0), 1))
        cos_part = torch.cos(dec[:, :self.obj_dim - 1] * (torch.pi / 2))
        cum_prod = torch.fliplr(torch.cumprod(torch.cat((ones_vector, cos_part), dim=1), dim=1))
        sin_part = torch.sin(torch.fliplr(dec[:, :self.obj_dim - 1] * (torch.pi / 2)))
        pop_obj = (1 + g.view(-1, 1)) * cum_prod * torch.cat((ones_vector, sin_part), dim=1)
        return pop_obj

    def get_optimal_solutions(self, size):
        return super().get_optimal_solutions1(size)
